import torch
from torch import nn


class RMSNorm(nn.Module):
    """
    这段代码定义了 RMSNorm（Root Mean Square Normalization，均方根归一化） 模块，
    这是一种在大语言模型（如 LLaMA、GPT 系列）中广泛使用的归一化技术，核心作用是稳定模型训练过程中的梯度，加速收敛，
    同时相比传统的 LayerNorm 减少了计算量（无需减去均值）。以下从模块功能、核心方法、细节优化等方面展开分析：
    一、核心背景：RMSNorm 是什么？
    RMSNorm 是 LayerNorm 的简化变体，其核心思想是仅通过 “均方根” 对特征进行归一化，而非 LayerNorm 的 “减均值 + 除标准差”。具体公式为：
    其中：\(\gamma\) 是可学习的缩放参数（对应代码中的 self.weight）；\(\epsilon\) 是防止分母为 0 的微小值（代码中默认 1e-6）；
    相比 LayerNorm 少了 “减均值” 步骤，计算更高效，且在大模型中常表现出更好的训练稳定性。


    模块通过 forward 方法分发逻辑，根据是否传入 residual（残差连接张量），调用两个不同的归一化方法：rms_forward（纯归一化）和 add_rms_forward（归一化 + 残差连接融合）。
    """
    def __init__(
        self,
        hidden_size: int,  # 输入特征的维度（如 Transformer 中的 hidden_dim）
        eps: float = 1e-6, # 防止分母为 0 的微小值
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(hidden_size)) # 可学习的缩放参数：初始化为全 1，形状与输入特征维度一致

    @torch.compile # 启用 PyTorch 编译优化，提升执行速度
    def rms_forward(
        self,
        x: torch.Tensor,
    ) -> torch.Tensor:
        """
        核心逻辑：严格遵循 RMSNorm 公式，分 3 步：
① 数据类型转换（低精度 → float32）→ ② 计算均方根并归一化 → ③ 恢复类型 + 缩放。
细节优化：
torch.float32 计算：低精度（如 float16）计算均方根时易出现数值溢出或精度损失，转为 float32 可提升稳定性，最后恢复原始类型不影响模型整体精度。
原地操作（mul_）：用 x.mul_(...) 替代 x = x * ...，减少中间张量创建，节省 GPU 内存。
keepdim=True：确保 var 最后一维为 1，与输入 x 形状（[batch, seq_len, hidden_size]）可广播（避免手动 reshape，简化代码）。
        :param x:
        :return:
        """
        orig_dtype = x.dtype # 保存输入原始数据类型（如 float16/bfloat16）
        x = x.to(torch.float32)  # 转换为 float32 计算，避免低精度数值不稳定
        var = x.pow(2).mean(dim=-1, keepdim=True) # 1. 计算均方根：先平方 → 最后一维求均值 → 开根号（+eps 防除零） # 形状：[batch, seq_len, 1]（假设输入是 [batch, seq_len, hidden_size]）
        x.mul_(torch.rsqrt(var + self.eps)) # x = x / sqrt(var + eps)（原地操作，节省内存）
        x = x.to(orig_dtype).mul_(self.weight) # 2. 恢复原始数据类型 + 应用可学习缩放 # 与 self.weight 逐元素相乘（广播机制）
        return x

    @torch.compile
    def add_rms_forward(
        self,
        x: torch.Tensor,
        residual: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        核心场景：在 Transformer 架构中，归一化常与残差连接配合（如 LayerNorm(x + Attention(x))）。此方法将 “残差相加” 与 “归一化” 融合，减少中间计算步骤。
关键区别：
先将 x 与 residual 相加（融合残差），再进行 RMS 归一化；
返回两个值：归一化后的特征 x + 融合后的残差 residual（供下一层使用）。
作用：避免单独计算残差连接的中间张量，提升计算效率并节省内存（尤其在大模型长序列场景）。
        :param x:
        :param residual:
        :return:
        """
        orig_dtype = x.dtype # 1. 融合“残差连接 + 数据类型转换”：x = x（float32） + residual（float32）
        x = x.to(torch.float32).add_(residual.to(torch.float32))
        residual = x.to(orig_dtype) # 保存融合后的残差（用于后续层的残差连接）
        var = x.pow(2).mean(dim=-1, keepdim=True) # 2. 与 rms_forward 一致的归一化逻辑
        x.mul_(torch.rsqrt(var + self.eps))
        x = x.to(orig_dtype).mul_(self.weight) # 3. 恢复类型 + 缩放
        return x, residual

    def forward(
        self,
        x: torch.Tensor,
        residual: torch.Tensor | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        """
        灵活性：支持两种使用场景：
仅需归一化（如模型某层无残差连接）：调用 rms_forward，返回单一张量；
归一化 + 残差融合（如 Transformer 层）：调用 add_rms_forward，返回 tuple（符合后续层的残差输入需求）。
        :param x:
        :param residual:
        :return:
        """
        if residual is None:
            return self.rms_forward(x) # 无残差：仅返回归一化结果
        else:
            return self.add_rms_forward(x, residual) # 有残差：返回（归一化结果, 融合后残差）
